package org.apache.sad.core.framework.jaxrpc.handler;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Calendar;
import java.util.TimeZone;
import java.util.Iterator;

import javax.xml.namespace.QName;
import javax.xml.soap.Name;
import javax.xml.rpc.handler.GenericHandler;
import javax.xml.rpc.handler.HandlerInfo;
import javax.xml.rpc.handler.MessageContext;
import javax.xml.rpc.handler.soap.SOAPMessageContext;
import javax.xml.soap.AttachmentPart;
import javax.xml.soap.SOAPElement;
import javax.xml.soap.SOAPEnvelope;
import javax.xml.soap.SOAPHeader;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPHeaderElement;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * The client handler is used to create WS-Security header and 
 * propagate service context when a service is calling another service 
 * via webservice interface.
 */
public class WsseClientHandler extends GenericHandler {
	public final static String SECURITY_ELEMENT = "Security";
	public final static String SECURITY_PREFIX = "wsse";
	public final static String SECURITY_NS = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd";
	public final static String MUST_UNDERSTAND_ATTRIB = "mustUnderstand";
	public final static String MUST_UNDERSTAND_ATTRIB_VALUE = "1";
	public final static String UNT_ELEMENT = "UsernameToken";
	public final static String WSU_PREFIX = "wsu";
	public final static String WSU_NS = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";
	public final static String USERNAME_ELEMENT = "Username";	
	public final static String SEPARATOR = ":";
	public final static String PASSWORD_ELEMENT = "Password";
	public final static String TYPE_ATTRIB = "Type";
	public final static String PASSWORD_TYPE_TEXT = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-username-token-profile-1.0#PasswordText";
	public final static String PASSWORD_TYPE_DIGEST = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-username-token-profile-1.0#PasswordDigest";
	public final static String NONCE_ELEMENT = "Nonce";
	public final static String CREATED_ELEMENT = "Created";
	public final static String ENCODING_TYPE_ATTRIB = "EncodingType";
	public final static String ENCODING_TYPE_ATTRIB_VALUE = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-soap-message-security-1.0#Base64Binary";
	public final static String TIMESTAMP_ELEMENT = "Timestamp";
	public final static String EXPIRES_ELEMENT = "Expires";
	
	private static Log logger = LogFactory.getLog(WsseClientHandler.class);
	protected HandlerInfo handlerInfo;
	private SecureRandom random;
	private boolean addTimestamp;
	private boolean usePasswordDigest;
	private boolean addNonce;
	private boolean wsse;
	
	public void init(HandlerInfo hi) {
	    handlerInfo = hi;
	    try {
	        random = SecureRandom.getInstance("SHA1PRNG");
	    } 
	    catch (NoSuchAlgorithmException nsae) {
	        logger.warn("Error creating random generator", nsae);
	        throw new RuntimeException(nsae);
	    }
	    wsse=false;
	    addNonce = false;
        usePasswordDigest = false;
        addTimestamp = false;
        
        Boolean b = (Boolean)handlerInfo.getHandlerConfig().get("addNonce");        
	    if (b != null)
	    	addNonce = b.booleanValue();
	    
	    b = (Boolean)handlerInfo.getHandlerConfig().get("wsse");
	    if (b != null)
	    	wsse = b.booleanValue();
	    
	    b = (Boolean)handlerInfo.getHandlerConfig().get("usePasswordDigest");
	    if (b != null)
	    	usePasswordDigest = b.booleanValue();
		
	    b = (Boolean)handlerInfo.getHandlerConfig().get("addTimestamp");
	    if (b != null)
	    	addTimestamp = b.booleanValue();	
	}

    public boolean handleRequest(MessageContext msgContext) {        
        SOAPMessage message = ((SOAPMessageContext)msgContext).getMessage();
        try
        {
        	//Add WS-Security header to the SOAP Message If wsse is enabled
        	if(wsse){
        		SOAPHeader header = message.getSOAPPart().getEnvelope().getHeader();
            	SOAPEnvelope envelope = message.getSOAPPart().getEnvelope();
            	String envNS = envelope.getPrefix();
            	SOAPElement secHeader = header.addChildElement(SECURITY_ELEMENT, SECURITY_PREFIX, SECURITY_NS);
            	secHeader.addAttribute(envelope.createName(envNS + SEPARATOR + MUST_UNDERSTAND_ATTRIB), MUST_UNDERSTAND_ATTRIB_VALUE);
            	
            	Calendar rightNow = Calendar.getInstance();        	
            	String createdTS = createTS(rightNow);
            	
            	createTimestampElement(secHeader, rightNow, createdTS);
            	        	
            	createUNT(secHeader, createdTS, envelope.createName(TYPE_ATTRIB));
        	}
        	
        		
        	return super.handleRequest(msgContext);            
        }
        catch (Exception e) {
        	logger.warn("Error creating WS-Security Header", e);
        	throw new RuntimeException(e);
        }
    }
    
    public boolean handleResponse(MessageContext msgContext) {
    	SOAPMessage message = ((SOAPMessageContext)msgContext).getMessage();
        try
        {
        	//Process Headers in the SOAP Message
        	SOAPEnvelope envelope = message.getSOAPPart().getEnvelope();
        	SOAPHeader header = envelope.getHeader();
        	
        	if (header != null) {
        		Iterator it = header.examineAllHeaderElements();
        		while (it.hasNext()) {
        			SOAPHeaderElement el = (SOAPHeaderElement)it.next();
        			if (el.getMustUnderstand())
        				el.setMustUnderstand(false);
        		}
        	}
        	this.processInboundAttachments((SOAPMessageContext)msgContext);
        }
        catch (Exception e) {
        	logger.warn("Error processing Soap Response", e);
        	throw new RuntimeException(e);
        }
    	return super.handleResponse(msgContext);
    }
    
    /**
     * Process any attachments in the SOAP message
     *
     * @param messageContext
     *            SOAPMessageContext the soap message context
     */
    private void processInboundAttachments(SOAPMessageContext messageContext) {
        Iterator attIter = messageContext.getMessage().getAttachments(); 
        
        // Check if there are any attachments
        if (attIter == null || !attIter.hasNext())
            return;

        try {
            // iterate through attachments
            AttachmentPart att;

            while (attIter.hasNext()) {
                // get the next attachment
                att = (AttachmentPart) attIter.next();
                logger.info("Received attachment of type: " + att.getContentType());
                //ServiceContextHandler.getInstance().getServiceContext().addOutboundAttachments(new SoapAttachment(att));
            }
        } catch (Exception e) {
            logger.warn("Error processing attachments", e);
        } finally {
            messageContext.getMessage().removeAllAttachments();
        }
    }
    
    private void createUNT(SOAPElement secHeader, String createdTS, Name typeAttrib)
    	throws SOAPException, Exception
    {
        SOAPElement node;        
    	SOAPElement unt = secHeader.addChildElement(UNT_ELEMENT, SECURITY_PREFIX);        	
    	node = unt.addChildElement(USERNAME_ELEMENT, SECURITY_PREFIX);
    	node.addTextNode((String)handlerInfo.getHandlerConfig().get("username"));
    	node = unt.addChildElement(PASSWORD_ELEMENT, SECURITY_PREFIX);
    	
    	String pwdType = PASSWORD_TYPE_TEXT;
    	String pwd = (String)handlerInfo.getHandlerConfig().get("password");
    	String nonce = null;
    	        	
    	if (usePasswordDigest || addNonce) {
    		nonce = createNonce();            	
    	}
    	
    	if (usePasswordDigest) {
    		pwdType = PASSWORD_TYPE_DIGEST;
    		pwd = this.createPasswordDigest(pwd, nonce, createdTS);
    	}
    	    	
    	node.addAttribute(typeAttrib, pwdType);        	        	
        node.addTextNode(pwd);            	
        
        if (usePasswordDigest || addNonce) {                    
        	unt.addNamespaceDeclaration(WSU_PREFIX, WSU_NS);        
        	node = unt.addChildElement(NONCE_ELEMENT, SECURITY_PREFIX);
        	//node.addAttribute(envelope.createName(ENCODING_TYPE_ATTRIB), ENCODING_TYPE_ATTRIB_VALUE);
        	node.addTextNode(nonce);    
    		node = unt.addChildElement(CREATED_ELEMENT, WSU_PREFIX);
    		node.addTextNode(createdTS);
        }        	
    
    }
    
    private void createTimestampElement(SOAPElement secHeader, Calendar rightNow, String createdTS) 
    	throws SOAPException 
    {
    	if (addTimestamp) {
    		rightNow.add(Calendar.MINUTE, 1);
    		String expiresTS = createTS(rightNow);
    		
    		SOAPElement tsElement = secHeader.addChildElement(TIMESTAMP_ELEMENT, WSU_PREFIX, WSU_NS);
    		SOAPElement node = tsElement.addChildElement(CREATED_ELEMENT, WSU_PREFIX);
    		node.addTextNode(createdTS);
    		
    		node = tsElement.addChildElement(EXPIRES_ELEMENT, WSU_PREFIX);
    		node.addTextNode(expiresTS);
    	}
    }
    
    public QName[] getHeaders() {
        return handlerInfo.getHeaders();
    }    
        		
    /**
     * Creates and adds a Nonce element to this UsernameToken
     */
    protected String createNonce() {
    	byte[] nonceValue = new byte[16];
        random.nextBytes(nonceValue);
        Base64 encoder = new Base64();
        return new String(encoder.encode(nonceValue));      
    }

    /**
     * Creates a timestamp to be used by wsu:Created and wsu:Expires
     */
    protected String createTS(Calendar rightNow) {
    	DateFormat zulu = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
        zulu.setTimeZone(TimeZone.getTimeZone("UTC"));
                  
        return zulu.format(rightNow.getTime());      
    }       
    
    protected String createPasswordDigest(String pwd, String nonce, String createdTS) throws Exception {    	
    	Base64 encoder = new Base64();
    	
    	MessageDigest sha = MessageDigest.getInstance("SHA-1");
    	sha.reset();
    	
    	byte [] b1 = encoder.decode(nonce.getBytes());
    	byte [] b2 = createdTS.getBytes("UTF-8");
    	byte [] b3 = pwd.getBytes("UTF-8");
    	byte [] b4 = new byte[b1.length + b2.length + b3.length];
    	
    	int offset = 0;
    	System.arraycopy(b1, 0, b4, offset, b1.length);
    	offset += b1.length;
    	
    	System.arraycopy(b2, 0, b4, offset, b2.length);
    	offset += b2.length;
    	
    	System.arraycopy(b3, 0, b4, offset, b3.length);
    	
    	sha.update(b4);
    	String pwdDigest = new String(encoder.encode(sha.digest()));
    	return pwdDigest;    	    	
    }
}